#include <chrono>
#include "superpoint.h"
#include <onnxruntime_cxx_api.h>

SuperPoint::SuperPoint(const std::string sp_onnx_path, int num_threads){
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "SuperPoint");
	Ort::SessionOptions sessionOptions;
	OrtCUDAProviderOptions options;
 	options.device_id = 0, 1, 2, 3;
	sessionOptions.AppendExecutionProvider_CUDA(options);
	sessionOptions.SetIntraOpNumThreads(num_threads);
	sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    sp_session_ = std::shared_ptr<Ort::Session>(new Ort::Session(env, sp_onnx_path.c_str(), sessionOptions));
}

void SuperPoint::DetectAndCompute(cv::Mat image, std::vector<cv::KeyPoint>& kpts, cv::Mat& desc){
    auto t1 = std::chrono::high_resolution_clock::now();
	Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "SuperPoint");
	auto t2 = std::chrono::high_resolution_clock::now();
	auto t12 = std::chrono::duration_cast<std::chrono::microseconds>(t2-t1);
	std::cout<<"Creat env time is: "<<t12.count()/1000<<" ms"<<std::endl;
	// 1.pre processing
    cv::Mat gray_img;
    if(image.channels() != 1){
        cv::cvtColor(image, gray_img, cv::COLOR_BGR2GRAY);
    }
    else{
        gray_img = image;
    }
    std::vector<float> img_data = ApplyTransform(gray_img);
    std::vector<int64_t> input_shape{1, 1, gray_img.rows, gray_img.cols};
	auto t3 = std::chrono::high_resolution_clock::now();
	auto t23 = std::chrono::duration_cast<std::chrono::microseconds>(t3-t2);
	std::cout<<"Before infer time is: "<<t23.count()/1000<<" ms"<<std::endl;
    // 2.infer
    Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
	Ort::Value inputTensor = Ort::Value::CreateTensor<float>(memoryInfo, img_data.data(), img_data.size(), input_shape.data(), input_shape.size());
    const char* input_names[] = {"image"};
	const char* output_names[] = {"keypoints", "scores", "descriptors"};
    Ort::RunOptions run_options;
	std::vector<Ort::Value> outputs = sp_session_->Run(run_options, input_names, &inputTensor, 1, output_names, 3);
	auto t4 = std::chrono::high_resolution_clock::now();
	auto t34 = std::chrono::duration_cast<std::chrono::microseconds>(t4-t3);
	std::cout<<"Run infer time is: "<<t34.count()/1000<<" ms"<<std::endl;

    // 3.parse output
    std::vector<int64_t> kpshape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
	int64* kp = (int64*)outputs[0].GetTensorMutableData<void>();
	int keypntcounts = kpshape[1];
	kpts.resize(keypntcounts);
    for (int i = 0; i < keypntcounts; i++)
	{
		cv::KeyPoint p;
		int index = i * 2;
		p.pt.x = (float)kp[index];
		p.pt.y = (float)kp[index + 1];
		kpts[i] = p;
	}

    std::vector<int64_t> desshape = outputs[2].GetTensorTypeAndShapeInfo().GetShape();
	float* des = (float*)outputs[2].GetTensorMutableData<void>();

	cv::Mat desmat;
	desmat.create(cv::Size(desshape[2], desshape[1]), CV_32FC1);
	for (int h = 0; h < desshape[1]; h++)
	{
		for (int w = 0; w < desshape[2]; w++)
		{
			int index = h * desshape[2] + w;
			desmat.at<float>(h, w) = des[index];
		}
	}
	desmat.copyTo(desc);
	auto t5 = std::chrono::high_resolution_clock::now();
	auto t45 = std::chrono::duration_cast<std::chrono::microseconds>(t4-t3);
	std::cout<<"After infer time is: "<<t45.count()/1000<<" ms"<<std::endl;
}

std::vector<float> SuperPoint::ApplyTransform(const cv::Mat image){
    cv::Mat resized, floatImage;
	image.convertTo(floatImage, CV_32FC1);

	vector<float> imgData;
	for (int h = 0; h < image.rows; h++)
	{
		for (int w = 0; w < image.cols; w++)
		{
			imgData.push_back(floatImage.at<float>(h, w) / 255.0f);
		}
	}
	return imgData;
}

void SuperPoint::vis(const cv::Mat image, const std::vector<cv::KeyPoint> kpts){
    cv::Mat img_with_kps;
    cv::drawKeypoints(image, kpts, img_with_kps);
    cv::imshow("SuperPoint Detect Result!!!", img_with_kps);
    cv::waitKey(0);
}